# -*- coding: UTF-8 -*-
"""
    @Author:YTQ
    @Time: 2022/11/28 17:34
    Description:
    
"""

import torch
from utils.utils import get_label, evaluate
from DataLoad.NewsDataSet import DataSet
from torch.utils import data
import torch.nn as nn
from modle.TextCNN import TextCNN
import logging


def train(config):
    # 类别
    # label, label_id = get_label()
    # 训练数据集
    train_dataset = DataSet(types='train')
    train_loader = data.DataLoader(train_dataset, batch_size=128, shuffle=True)
    # 验证数据集
    dev_dataset = DataSet(types='dev')
    dev_loader = data.DataLoader(dev_dataset, batch_size=128, shuffle=True)
    # 加载模型
    model = TextCNN(config).to(config.device)
    if not config.isUseGpu:
        logging.warning('当前使用CPU训练数据，训练速度较慢，切换至GPU平台高速训练')
    # 优化器
    optimizer = torch.optim.Adam(model.parameters(), lr=config.learning_rate)
    # 损失
    loss_fn = nn.CrossEntropyLoss()
    #
    for i in range(config.epochs_num):
        for b, (x, mask, target) in enumerate(train_loader):
            # 使用了GPU
            if config.isUseGpu:
                x = x.to(config.device)
                mask = mask.to(config.device)
                target = target.to(config.device)

            pred = model(x, mask)
            loss = loss_fn(pred, target)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            # 限制
            if b % 50 != 0:
                continue
            y_pred = torch.argmax(pred, dim=1)

            if config.isUseGpu:
                report = evaluate(y_pred.cpu().data.numpy(), target.cpu().data.numpy(), output_dic=True)
            else:
                report = evaluate(y_pred.data.numpy(), target.data.numpy(), output_dic=True)

            with torch.no_grad():
                dev_input, dev_mask, dev_target = iter(dev_loader).__next__()
                if config.isUseGpu:
                    dev_input = dev_input.to(config.device)
                    dev_mask = dev_mask.to(config.device)
                    dev_target = dev_target.to(config.device)
                dev_pred = model(dev_input, dev_mask)
                dev_pred_ = torch.argmax(dev_pred, dim=1)
                if config.isUseGpu:
                    dev_report = evaluate(dev_pred_.cpu().data.numpy(), dev_target.cpu().data.numpy(), output_dic=True)
                else:
                    dev_report = evaluate(dev_pred_.data.numpy(), dev_target.data.numpy(), output_dic=True)

            logging.info(f'epoch:{i}, batch:{b}, loss:{round(loss.item(), 5)}, train_acc:{report["accuracy"]}, dev_acc{dev_report["accuracy"]}')

        # save
        torch.save(model.state_dict(), f'{config.save_path}{i}.pth')

    #
    logging.info(f'训练完毕，模型存储路径：{config.save_path}')
    logging.info('赶快使用test.py检测模型的准确度')
